import torch
import json
# 替换 torchsummary 为 torchinfo
from torchinfo import summary  
from models.price_predictor import PricePredictor
from datasets.dataset import TicketDataset

def print_model_structure():
    CONFIG_PATH = "movie-ticket-bidding/config/feature_attention.json"
    DATA_PATH = "movie-ticket-bidding/data/processed/test.pkl"
    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
    
    print(f"使用设备: {DEVICE}")
    if DEVICE == "cuda":
        print(f"CUDA设备: {torch.cuda.get_device_name(0)}")
    
    # 获取输入维度
    dataset = TicketDataset(DATA_PATH)
    input_dim = dataset[0][0].shape[0]
    print(f"自动获取输入特征维度: {input_dim}\n")
    
    # 加载模型配置
    with open(CONFIG_PATH, "r") as f:
        model_cfg = json.load(f)["model"]
        print(f"已加载模型配置: {CONFIG_PATH}")
        print(f"模型配置详情: {json.dumps(model_cfg, indent=2)}\n")
    
    # 创建模型
    model = PricePredictor(CONFIG_PATH, input_dim).to(DEVICE)
    model.eval()
    
    # 打印整体结构
    print("="*80)
    print("模型整体结构:")
    print(model)
    print("="*80 + "\n")
    
    # 打印每一层输出形状（验证输入维度正确性）
    print("各层输出形状:")
    with torch.no_grad():
        test_input = torch.randn(2, input_dim).to(DEVICE)
        print(f"输入形状: {test_input.shape}")
        x = test_input
        for i, layer in enumerate(model.layers):
            x = layer(x)
            print(f"层 {i+1} ({layer.__class__.__name__}) 输出形状: {x.shape}")
    
    # 使用 torchinfo 打印参数详情（兼容 Transformer 结构）
    print("\n" + "="*80)
    print("模型参数详情（torchinfo）:")
    summary(
        model,
        input_size=(input_dim,),  # 单个样本形状（不含batch），模型期望 (batch, 7)
        batch_size=2,  # 显式指定batch_size，可选
        device=DEVICE,
        col_names=["input_size", "output_size", "num_params", "trainable"],  # 显示的列
        col_width=20,
        row_settings=["var_names"]
    )
    print("="*80)
    
    return model

if __name__ == "__main__":
    print_model_structure()